from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader
from langchain.embeddings import OpenAIEmbeddings 
from langchain.indexes import VectorstoreIndexCreator
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.indexes import VectorstoreIndexCreator
filename1 = "state_of_the_union.txt"
filename2 = "story of meici.txt"

def openai_with_eng(filename):
    embedding = OpenAIEmbeddings(model="text-embedding-ada-002",openai_api_base="http://localhost:8000/v1",openai_api_key="EMPTY")
    loader = TextLoader(filename,autodetect_encoding=True)
    index = VectorstoreIndexCreator(embedding=embedding).from_loaders([loader])
    llm = ChatOpenAI(model="gpt-3.5-turbo",openai_api_base="http://localhost:8000/v1",openai_api_key="EMPTY")

    questions = [
        "Who is the speaker",
        "What did the president say about Ketanji Brown Jackson",
        "What are the threats to America",
        "Who are mentioned in the speech",
        "Who is the vice president",
        "How many projects were announced",
    ]

    for query in questions:
        print("Query:", query)
        print("Answer:", index.query(query, llm=llm))

def openai_with_cn(filename):
    # 改成中文分词
    embedding = HuggingFaceEmbeddings(model_name='shibing624/text2vec-base-chinese')
    loader = TextLoader(filename,autodetect_encoding=True)
    index = VectorstoreIndexCreator(embedding=embedding).from_loaders([loader])
    llm = ChatOpenAI(model="gpt-3.5-turbo",openai_api_base="http://localhost:8000/v1",openai_api_key="EMPTY")

    questions = [
        "Who is the speaker",
        "What did the president say about Ketanji Brown Jackson",
        "What are the threats to America",
        "Who are mentioned in the speech",
        "Who is the vice president",
        "How many projects were announced",
    ]

    # questions_cn = [
    #     "故事的主角是谁",
    #     "故事中出现多少个男性角色",
    #     "故事中出现多少个女性角色",
    #     "故事中出现多少个人物",
    #     "故事中出现多少个动物",
    #     "故事中最大的官职是什么",
    # ]

    for query in questions:
        print("Query:", query)
        print("Answer:", index.query(query, llm=llm))

def openai_with_cn_question(filename):
    # 改成中文分词,中文文章以及中文问题，会有token超过限制的问题
    embedding = HuggingFaceEmbeddings(model_name='shibing624/text2vec-base-chinese')
    loader = TextLoader(filename,autodetect_encoding=True)
    index = VectorstoreIndexCreator(embedding=embedding).from_loaders([loader])
    llm = ChatOpenAI(model="gpt-3.5-turbo",openai_api_base="http://localhost:8000/v1",openai_api_key="EMPTY")

    questions = [
        "故事的主角是谁",
        "故事中出现多少个男性角色",
        "故事中出现多少个女性角色",
        "故事中出现多少个人物",
        "故事中出现多少个动物",
        "故事中最大的官职是什么",
    ]

    for query in questions:
        print("Query:", query)
        print("Answer:", index.query(query, llm=llm))


# openai_with_eng(filename1)
# openai_with_cn(filename1)
openai_with_cn_question(filename2)